//
// Created by hjie on 23-5-31.
//

#include "message_attr.h"
#include <arpa/inet.h>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <endian.h>
#include <strings.h>

#include <zlib.h>

#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/sha.h>

MessageAttr::MessageAttr(MessageAttrType type)
{
    m_attr_type = type;
}

MessageAttr::~MessageAttr()
{

}

std::shared_ptr<MessageAttr> MessageAttr::ProcessMessageAttr(const uint8_t *data, int length, int &use_length)
{
    uint16_t attr_type = 0;
    memcpy((uint8_t*)&attr_type, data, 2);
    attr_type = be16toh(attr_type);

    MessageAttr * attr = nullptr;
    switch (attr_type)
    {
        case EmXorPeerAddress:
        case EmXorMappedAddress:
        case EmXorRelayedAddress:
        {
            attr = new AttrXorMappedAddress();
            attr->m_attr_type = (MessageAttrType)attr_type;
            break;
        }
        case EmRequestedTransport:
        {
            attr = new AttrRequestedTransport();
            break;
        }
        case EmErrorCode:
        {
            attr = new AttrErrorCode();
            break;
        }
        case EmUserName:
        {
            attr = new AttrUserName();
            break;
        };
        case EmNONCE:
        {
            attr = new AttrNonce();
            break;
        }
        case EmRealm:
        {
            attr = new AttrRealm();
            break;
        }
        case EmLifeTime:
        {
            attr = new AttrLifeTime();
            break;
        }
        case EmData:
        {
            attr = new AttrData();
            break;
        }
        case EmFingerPrint:
        {
            attr = new AttrCrc32();
            break;
        }
        case EmMessageIntegrity:
        {
            attr = new AttrMessageIntegrity();
            break;
        }
        default:
        {
            uint16_t attr_length = 0;
            memcpy(&attr_length, data + 2, 2);
            attr_length = be16toh(attr_length);
            printf("drop stun message attr : %d and %4x.\n", attr_type, attr_type);

            use_length += 4 + attr_length;
            break;
        }
    }
    if (attr)
    {
        uint16_t attr_length = 0;
        memcpy(&attr_length, data + 2, 2);
        attr_length = be16toh(attr_length);
        if (attr_length <= length - 4)
        {
            attr->m_length = attr_length;
            use_length += 4;
            bool result = attr->ParseAttr(data + 4, attr_length, use_length);
            if (!result)
            {
                delete attr;
                attr = nullptr;
                use_length = -1;
            }
        }
        else
        {
            delete attr;
            attr = nullptr;
            // drop stun packet
            use_length = -1;
        }
    }
    return attr == nullptr ? nullptr : std::shared_ptr<MessageAttr>(attr);
}

void MessageAttr::ComposeBase(WriteBuffer &writer)
{
    writer.Write(htobe16(m_attr_type));
    writer.Write(htobe16(m_length));
}

void MessageAttr::PaddingData(WriteBuffer &writer, int padding_size)
{
    for (int i = 0; i < padding_size; i++)
    {
        writer.Write(uint8_t(0));
    }
}

AttrXorMappedAddress::AttrXorMappedAddress(): MessageAttr(EmXorMappedAddress)
{
    m_length = 8;
    // EmXor
}

bool AttrXorMappedAddress::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    int index = 0;
    index++;
    m_protocol = data[index];
    index++;
    memcpy(&m_port, data + index, 2);
    index += 2;
    memcpy(&m_addr, data + index, 4);
    index += 4;
    use_length += m_length;
    return true;
}

void AttrXorMappedAddress::ParseAddrInfo(int protocol_type, const struct sockaddr *addr, int length)
{
    // access use localxxx 
    switch (protocol_type)
    {
        case IPPROTO_IP:
        {
            if (sizeof(struct sockaddr_in) == length)
            {
                const struct sockaddr_in * address = (const struct sockaddr_in*)addr;
                m_port = address->sin_port;
                m_addr = address->sin_addr.s_addr;
                m_protocol = 0x01;
            }
            break;
        }
        case IPPROTO_IPV6:
        {
            if (sizeof(struct sockaddr_in6) == length)
            {

            }
            break;
        }
        default:
            break;
    }
}

std::string AttrXorMappedAddress::LocalIp()
{
    char ip[128] = {0};
    inet_ntop(AF_INET, &m_addr, ip, sizeof(ip));
    return ip;
}

uint16_t AttrXorMappedAddress::LocalPort()
{
    return ntohs(m_port);
}

uint16_t AttrXorMappedAddress::DecoderPort()
{
    uint16_t port = ntohs(m_port);
    return port ^ (magic_cookie >> 16);
}

uint16_t AttrXorMappedAddress::EncoderPort()
{
    uint16_t port = ntohs(m_port) ^ (magic_cookie >> 16);
    return htons(port);
}

uint32_t AttrXorMappedAddress::EncoderAddr()
{
    uint32_t addr = m_addr ^ htobe32(magic_cookie);
    return addr;
}

std::string AttrXorMappedAddress::DecoderAddr()
{
    uint32_t addr = m_addr ^ be32toh(magic_cookie);
    char ip[128] = {0};
    inet_ntop(AF_INET, &addr, ip, sizeof(ip));
    return ip;
}

int AttrXorMappedAddress::GetAttrLength()
{
    // const value 4(type and attr length) + attr length
    return m_length + 4;
}

void AttrXorMappedAddress::ComposeAttr(WriteBuffer &writer)
{
    ComposeBase(writer);

    writer.Write((uint8_t)00);// padding 8 bit
    writer.Write(m_protocol);
    writer.Write(EncoderPort());
    writer.Write(EncoderAddr());
}

AttrRequestedTransport::AttrRequestedTransport() : MessageAttr(EmRequestedTransport)
{
    m_length = 4;
}

bool AttrRequestedTransport::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    uint8_t method = *data;
    use_length += 4;

    if (method == 0x11)
    {
        m_request_transport = "UDP";
    }
    else
    {
        m_request_transport = "TCP";
    }
    return true;
}

int AttrRequestedTransport::GetAttrLength()
{
    return m_length + 4;
}

void AttrRequestedTransport::ComposeAttr(WriteBuffer &writer)
{
    ComposeBase(writer);
    if (m_request_transport == "UDP")
    {
        writer.Write((uint8_t)0x11);
    }
    writer.Write(uint16_t(0));
    writer.Write(uint8_t(0));
}

void AttrRequestedTransport::SetTransport(const std::string &method)
{
    m_request_transport = method;
}


AttrRealm::AttrRealm() : MessageAttr(EmRealm)
{
    m_length = 0;
}

int AttrRealm::GetAttrLength()
{
    return m_length + 4;
}

void AttrRealm::SetRealm(const std::string &realm)
{
    m_realm  = realm;
    m_length = m_realm.size();
    if (m_length % 4 != 0)
    {
        m_length += (4 - m_length % 4);
    }
}

bool AttrRealm::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_realm.append((const char*)data, length);
    if (length % 4 != 0)
    {
        m_length = length + (4 - length % 4);
    }
    use_length += m_length;
    return true;
}

void AttrRealm::ComposeAttr(WriteBuffer &writer)
{
    uint16_t size = m_realm.size();
    m_length = size;
    ComposeBase(writer);
    writer.Write((const uint8_t*)m_realm.c_str(), size);
    if (size % 4 != 0)
    {
        uint8_t padding = 4 - (size % 4);
        for (int i = 0; i < padding; i++)
        {
            writer.Write(uint8_t(0));
        }
    }
}

/////////////////////////////////////////Error Code Attr////////////////////////////////////////////////////////////////
AttrErrorCode::AttrErrorCode() : MessageAttr(EmErrorCode)
{
    m_error_class = 0;
    m_error_code  = 0;
    m_length      = 4;
}

int AttrErrorCode::GetAttrLength()
{
    return m_length + 4;
}

bool AttrErrorCode::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    int index = 2;
    m_error_class = *(data + index);
    index++;
    m_error_code  = *(data + index);
    use_length += 4;
    return true;
}

void AttrErrorCode::ComposeAttr(WriteBuffer &writer)
{
    ComposeBase(writer);
    writer.Write((uint16_t)0000);
    writer.Write(m_error_class);
    writer.Write(m_error_code);
}

void AttrErrorCode::SetError(uint8_t error_class, uint8_t error_code)
{
    m_error_class = error_class;
    m_error_code  = error_code;
}

/////////////////////////////////////////Error Code Attr////////////////////////////////////////////////////////////////
AttrNonce::AttrNonce() : MessageAttr(EmNONCE)
{
    m_length = 32;
}

bool AttrNonce::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_nonce_data.append((const char *)data, length);
    use_length += length;
    return true;
}

void AttrNonce::ComposeAttr(WriteBuffer &writer)
{
    ComposeBase(writer);
    writer.Write((const uint8_t*)m_nonce_data.c_str(), 32);
}

void AttrNonce::SetNonce(const std::string & content)
{
    m_nonce_data = content;
    m_length = 32;
}

int AttrNonce::GetAttrLength()
{
    return 32 + 4;
}
/////////////////////////////////////////User Name Attr/////////////////////////////////////////////////////////////////
AttrUserName::AttrUserName() : MessageAttr(EmUserName)
{

}

bool AttrUserName::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_user_name.append((const char *)data, length);
    if (length % 4 != 0)
    {
        m_length = length + (4 - length % 4);
    }
    use_length += m_length;
    return true;
}

void AttrUserName::ComposeAttr(WriteBuffer &writer)
{
    int attr_length = m_length;
    m_length = m_user_name.size();
    ComposeBase(writer);
    writer.Write((const uint8_t*)m_user_name.c_str(), m_user_name.size());
    if (m_user_name.size() != attr_length)
    {
        int padding_size = attr_length - m_user_name.size();
        for (int i = 0; i < padding_size; i++)
        {
            writer.Write((uint8_t)00);
        }
    }
}

int AttrUserName::GetAttrLength()
{
    return m_length + 4;
}

void AttrUserName::SetUserName(const std::string &user_name)
{
    m_user_name = user_name;
    m_length    = m_user_name.size();
    if (m_length % 4 != 0)
    {
        m_length += (4 - m_length % 4);
    }
    else 
    {
        m_length = user_name.size();
    }
}

///////////////////////////////////////// Attr/////////////////////////////////////////////////////////////////
AttrLifeTime::AttrLifeTime() : MessageAttr(EmLifeTime)
{
    m_length = 4;
    m_left_time = 600;
}

bool AttrLifeTime::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    memcpy((uint8_t*)&m_left_time, data, 4);
    m_left_time = be32toh(m_left_time);
    length = 4;
    use_length += 4;
    return true;
}

void AttrLifeTime::ComposeAttr(WriteBuffer &writer)
{
    ComposeBase(writer);
    uint32_t data = be32toh(m_left_time);
    writer.Write(data);
}

int AttrLifeTime::GetAttrLength()
{
    return m_length + 4;
}

int AttrLifeTime::GetLifeTime()
{
    return m_left_time;
}

///////////////////////////////////////// Attr//////////////////////////////////////////////////////////////////////////

AttrData::AttrData() : MessageAttr(EmData)
{
    m_length = 0;
}

AttrData::~AttrData()
{
    m_data_vec.clear();
    m_data_vec.shrink_to_fit();
}

bool AttrData::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_data_vec.resize(length);
    memcpy(&(*m_data_vec.begin()), data, length);
    if (length % 4 != 0)
    {
        m_length = length + (4 - length % 4);
    }
    use_length += m_length;
    return true;
}

void AttrData::PushData(const uint8_t *data, int length)
{
    m_data_vec.resize(length);
    memcpy(&(*m_data_vec.begin()), data, length);
    if (length % 4 != 0)
    {
        m_length = length + (4 - length % 4);
    }
    else 
    {
        m_length = length;
    }
}

int AttrData::GetAttrLength()
{
    return m_length + 4;
}

void AttrData::ComposeAttr(WriteBuffer &writer)
{
    int size = m_length;
    m_length = m_data_vec.size();
    ComposeBase(writer);
    writer.Write(&*m_data_vec.begin(), m_data_vec.size());
    if (m_data_vec.size() != size)
    {
        int padding = size - m_data_vec.size();
        PaddingData(writer, padding);
    }
}


AttrCrc32::AttrCrc32():MessageAttr(EmFingerPrint)
{
    m_length = 4;
}

AttrCrc32::~AttrCrc32()
{

}

bool AttrCrc32::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_length = 4;
    memcpy((uint8_t*)&m_crc32_data, data, m_length);
    m_crc32_data = be32toh(m_crc32_data);
    // printf("crc 32 data : %8x and %8x.\n", m_crc32_data, m_crc32_data ^ 0x5354554E);
    use_length += 4;
    return true;
}

void AttrCrc32::ComposeAttr(WriteBuffer &writer)
{
    uint8_t * data = writer.Data();
    int length = writer.Length();
    uLong crc = crc32(0L, Z_NULL, 0);
    crc = crc32(crc, (const Bytef*)data, length);

    m_crc32_data = crc ^ finger_print;
    m_crc32_data = htobe32(m_crc32_data);

    ComposeBase(writer);
    writer.Write(m_crc32_data);
}

int AttrCrc32::GetAttrLength()
{
    return m_length;
}

bool AttrCrc32::CheckCrc32Data(const uint8_t *data, int length)
{
    uLong crc = crc32(0L, Z_NULL, 0);
    crc = crc32(crc, (const Bytef*)data, length);
    // printf("check crc 32 data : %8lx and %8lx.\n", crc, crc ^ finger_print);
    if (crc == (m_crc32_data ^ finger_print))
    {
        return true;
    }
    return false;
}

void AttrCrc32::ComputeCrc32(uint8_t *data, int length)
{

}

AttrMessageIntegrity::AttrMessageIntegrity():MessageAttr(EmMessageIntegrity)
{
    m_length = 20;
    bzero(m_hmac_data, sizeof(m_hmac_data));
}

AttrMessageIntegrity::~AttrMessageIntegrity()
{

}

int AttrMessageIntegrity::GetAttrLength()
{
    return m_length;
}

bool AttrMessageIntegrity::ParseAttr(const uint8_t *data, int length, int &use_length)
{
    m_length = 20;
    // unsigned char code_data[20] = {0};
    memcpy(m_hmac_data, data, m_length);
    // m_hmac_encode_value = (char*)code_data;
    use_length += 20;
    return true;
}

void AttrMessageIntegrity::ComposeAttr(WriteBuffer &writer)
{
    uint8_t * data = writer.Data();
    int length = writer.Length();
}

bool AttrMessageIntegrity::CheckHmacData(const uint8_t * data, int length)
{

    // unsigned char md_buffer[32] = {0};
    // unsigned int  md_length = 0;
    // HMAC(EVP_sha1(), (const unsigned char *)m_hmac_key.c_str(), m_hmac_key.size(), (const unsigned char *)data, length, md_buffer, &md_length);
    // compare md_buffer and m_hmac_encode_value
    return true;
}

bool AttrMessageIntegrity::CheckHmacData(const uint8_t * data, int length, const std::string & key)
{
    unsigned char md_buffer[32] = {0};
    unsigned int  md_length = 0;
    HMAC(EVP_sha1(), (const unsigned char *)key.c_str(), key.size(), (const unsigned char *)data, length, md_buffer, &md_length);
    if (md_length != m_length)
    {
        return false;
    }
    for (unsigned int index = 0; index < md_length; index++)
    {
        if (md_buffer[index] != m_hmac_data[index])
        {
            return false;
        }
    }
    return true;
}



